{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chat Prompts Customization"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prompt Setup\n",
    "\n",
    "Below, we take the default prompts and customize them to always answer, even if the context is not helpful."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms import ChatMessage, MessageRole\n",
    "from llama_index.prompts import ChatPromptTemplate\n",
    "\n",
    "# Text QA Prompt\n",
    "chat_text_qa_msgs = [\n",
    "    ChatMessage(\n",
    "        role=MessageRole.SYSTEM,\n",
    "        content=\"Always answer the question, even if the context isn't helpful.\",\n",
    "    ),\n",
    "    ChatMessage(\n",
    "        role=MessageRole.USER,\n",
    "        content=(\n",
    "            \"Context information is below.\\n\"\n",
    "            \"---------------------\\n\"\n",
    "            \"{context_str}\\n\"\n",
    "            \"---------------------\\n\"\n",
    "            \"Given the context information and not prior knowledge, \"\n",
    "            \"answer the question: {query_str}\\n\"\n",
    "        ),\n",
    "    ),\n",
    "]\n",
    "text_qa_template = ChatPromptTemplate(chat_text_qa_msgs)\n",
    "\n",
    "# Refine Prompt\n",
    "chat_refine_msgs = [\n",
    "    ChatMessage(\n",
    "        role=MessageRole.SYSTEM,\n",
    "        content=\"Always answer the question, even if the context isn't helpful.\",\n",
    "    ),\n",
    "    ChatMessage(\n",
    "        role=MessageRole.USER,\n",
    "        content=(\n",
    "            \"We have the opportunity to refine the original answer \"\n",
    "            \"(only if needed) with some more context below.\\n\"\n",
    "            \"------------\\n\"\n",
    "            \"{context_msg}\\n\"\n",
    "            \"------------\\n\"\n",
    "            \"Given the new context, refine the original answer to better \"\n",
    "            \"answer the question: {query_str}. \"\n",
    "            \"If the context isn't useful, output the original answer again.\\n\"\n",
    "            \"Original Answer: {existing_answer}\"\n",
    "        ),\n",
    "    ),\n",
    "]\n",
    "refine_template = ChatPromptTemplate(chat_refine_msgs)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using the Prompts\n",
    "\n",
    "Now, we use the prompts in an index query!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openai\n",
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"sk-...\"\n",
    "openai.api_key = os.environ[\"OPENAI_API_KEY\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index import VectorStoreIndex, SimpleDirectoryReader, ServiceContext\n",
    "from llama_index.llms import OpenAI\n",
    "\n",
    "documents = SimpleDirectoryReader(\"../../data/paul_graham/\").load_data()\n",
    "\n",
    "# Create an index using a chat model, so that we can use the chat prompts!\n",
    "service_context = ServiceContext.from_defaults(\n",
    "    llm=OpenAI(model=\"gpt-3.5-turbo\", temperature=0.1)\n",
    ")\n",
    "\n",
    "index = VectorStoreIndex.from_documents(documents, service_context=service_context)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Before Adding Templates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "I'm sorry, but the given context does not provide any information about Joe Biden.\n"
     ]
    }
   ],
   "source": [
    "print(index.as_query_engine().query(\"Who is Joe Biden?\"))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### After Adding Templates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Joe Biden is the 46th President of the United States.\n"
     ]
    }
   ],
   "source": [
    "print(\n",
    "    index.as_query_engine(\n",
    "        text_qa_template=text_qa_template, refine_template=refine_template\n",
    "    ).query(\"Who is Joe Biden?\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama-index",
   "language": "python",
   "name": "llama-index"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.0"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
